// 引入 Node.js 的 fs 模块，用于读取文件
const fs = require('fs');

// 数据预处理函数，将文本转换为适合 RNN 处理的格式
function preprocessText(text) {
  // 将文本转换为小写，然后使用正则表达式匹配所有单词，将结果存储在 words 数组中
  // /\w+/g 表示匹配一个或多个连续的字母、数字或下划线字符
  const words = text.toLowerCase().match(/\w+/g);
  // 使用 Set 数据结构去除 words 数组中的重复单词，再将其转换为数组，得到词汇表
  const vocabulary = [...new Set(words)];
  // 创建一个空对象 wordToIndex，用于将单词映射到对应的索引
  const wordToIndex = {};
  // 创建一个空对象 indexToWord，用于将索引映射到对应的单词
  const indexToWord = {};
  // 遍历词汇表中的每个单词和对应的索引
  vocabulary.forEach((word, index) => {
    // 将单词作为键，索引作为值，存入 wordToIndex 对象
    wordToIndex[word] = index;
    // 将索引作为键，单词作为值，存入 indexToWord 对象
    indexToWord[index] = word;
  });
  // 将 words 数组中的每个单词转换为对应的索引，得到输入序列
  const inputSequence = words.map(word => wordToIndex[word]);
  // 返回包含词汇表、单词到索引的映射、索引到单词的映射和输入序列的对象
  return { vocabulary, wordToIndex, indexToWord, inputSequence };
}

// 定义 RNN 类，用于实现循环神经网络
class RNN {
  // 构造函数，初始化 RNN 的输入大小、隐藏大小和输出大小，并初始化权重和偏置
  constructor(inputSize, hiddenSize, outputSize) {
    this.inputSize = inputSize;
    this.hiddenSize = hiddenSize;
    this.outputSize = outputSize;

    // 初始化输入到隐藏层的权重矩阵 Wxh，其形状为 (hiddenSize, inputSize)
    // 权重初始值在 [-0.05, 0.05) 范围内随机生成
    this.Wxh = this.randomMatrix(hiddenSize, inputSize);
    // 初始化隐藏层到隐藏层的权重矩阵 Whh，其形状为 (hiddenSize, hiddenSize)
    this.Whh = this.randomMatrix(hiddenSize, hiddenSize);
    // 初始化隐藏层到输出层的权重矩阵 Why，其形状为 (outputSize, hiddenSize)
    this.Why = this.randomMatrix(outputSize, hiddenSize);
    // 初始化隐藏层的偏置向量 bh，其长度为 hiddenSize，初始值都为 0
    this.bh = this.zeros(hiddenSize);
    // 初始化输出层的偏置向量 by，其长度为 outputSize，初始值都为 0
    this.by = this.zeros(outputSize);
  }

  // 生成指定行数和列数的随机矩阵，矩阵元素的值在 [-0.05, 0.05) 范围内
  randomMatrix(rows, cols) {
    const matrix = [];
    // 遍历每一行
    for (let i = 0; i < rows; i++) {
      const row = [];
      // 遍历每一列
      for (let j = 0; j < cols; j++) {
        // 生成一个在 [-0.05, 0.05) 范围内的随机数，并添加到当前行
        row.push(Math.random() * 0.1 - 0.05);
      }
      // 将当前行添加到矩阵中
      matrix.push(row);
    }
    return matrix;
  }

  // 生成指定长度的零向量
  zeros(size) {
    return new Array(size).fill(0);
  }

  // 激活函数：双曲正切函数 tanh
  // 将输入值 x 转换为其双曲正切值，用于引入非线性
  tanh(x) {
    return Math.tanh(x);
  }

  // tanh 函数的导数
  // 用于反向传播时计算梯度
  tanhDerivative(x) {
    return 1 - x * x;
  }

  // 前向传播函数，计算 RNN 在给定输入序列下的隐藏状态和输出
  forward(inputs) {
    // 存储每个时间步的隐藏状态
    const h = [];
    // 存储每个时间步的输出
    const y = [];
    // 初始化上一个时间步的隐藏状态为零向量
    let prevH = this.zeros(this.hiddenSize);

    // 遍历输入序列中的每个元素
    for (let t = 0; t < inputs.length; t++) {
      // 将当前输入转换为独热编码向量
      const x = this.oneHotEncode(inputs[t], this.inputSize);
      // 初始化当前时间步的隐藏状态为零向量
      const hNext = this.zeros(this.hiddenSize);
      // 计算当前时间步的隐藏状态
      for (let i = 0; i < this.hiddenSize; i++) {
        // 初始化当前隐藏单元的输入总和为偏置值
        let sum = this.bh[i];
        // 计算输入到隐藏层的加权和
        for (let j = 0; j < this.inputSize; j++) {
          sum += this.Wxh[i][j] * x[j];
        }
        // 计算上一个时间步的隐藏状态到当前隐藏状态的加权和
        for (let j = 0; j < this.hiddenSize; j++) {
          sum += this.Whh[i][j] * prevH[j];
        }
        // 对输入总和应用 tanh 激活函数，得到当前隐藏单元的值
        hNext[i] = this.tanh(sum);
      }
      // 初始化当前时间步的输出为零向量
      const yNext = this.zeros(this.outputSize);
      // 计算当前时间步的输出
      for (let i = 0; i < this.outputSize; i++) {
        // 初始化当前输出单元的输入总和为偏置值
        let sum = this.by[i];
        // 计算隐藏层到输出层的加权和
        for (let j = 0; j < this.hiddenSize; j++) {
          sum += this.Why[i][j] * hNext[j];
        }
        // 将输入总和作为当前输出单元的值
        yNext[i] = sum;
      }
      // 将当前时间步的隐藏状态添加到 h 数组中
      h.push(hNext);
      // 将当前时间步的输出添加到 y 数组中
      y.push(yNext);
      // 更新上一个时间步的隐藏状态为当前时间步的隐藏状态
      prevH = hNext;
    }

    // 返回包含每个时间步的隐藏状态和输出的对象
    return { h, y };
  }

  // 将指定索引转换为独热编码向量
  oneHotEncode(index, size) {
    // 生成指定长度的零向量
    const vector = this.zeros(size);
    // 将指定索引位置的值设为 1
    vector[index] = 1;
    return vector;
  }

  // 反向传播函数，计算权重和偏置的梯度
  backward(inputs, targets, h, y) {
    // 初始化输入到隐藏层的权重梯度矩阵，初始值都为 0
    const dWxh = this.zerosMatrix(this.hiddenSize, this.inputSize);
    // 初始化隐藏层到隐藏层的权重梯度矩阵，初始值都为 0
    const dWhh = this.zerosMatrix(this.hiddenSize, this.hiddenSize);
    // 初始化隐藏层到输出层的权重梯度矩阵，初始值都为 0
    const dWhy = this.zerosMatrix(this.outputSize, this.hiddenSize);
    // 初始化隐藏层的偏置梯度向量，初始值都为 0
    const dbh = this.zeros(this.hiddenSize);
    // 初始化输出层的偏置梯度向量，初始值都为 0
    const dby = this.zeros(this.outputSize);
    // 初始化下一个时间步的隐藏状态梯度为零向量
    let dhNext = this.zeros(this.hiddenSize);

    // 从最后一个时间步开始反向遍历输入序列
    for (let t = inputs.length - 1; t >= 0; t--) {
      // 初始化当前时间步的输出梯度向量，初始值都为 0
      const dy = this.zeros(this.outputSize);
      // 计算当前时间步的输出梯度
      for (let i = 0; i < this.outputSize; i++) {
        // 如果当前输出单元的索引等于目标索引，则梯度为预测值减 1，否则为预测值
        dy[i] = y[t][i] - (i === targets[t] ? 1 : 0);
      }
      // 计算隐藏层到输出层的权重梯度和输出层的偏置梯度
      for (let i = 0; i < this.outputSize; i++) {
        for (let j = 0; j < this.hiddenSize; j++) {
          // 累加当前时间步的隐藏层到输出层的权重梯度
          dWhy[i][j] += dy[i] * h[t][j];
        }
        // 累加当前时间步的输出层的偏置梯度
        dby[i] += dy[i];
      }
      // 初始化当前时间步的隐藏状态梯度向量，初始值都为 0
      const dh = this.zeros(this.hiddenSize);
      // 计算当前时间步的隐藏状态梯度
      for (let i = 0; i < this.hiddenSize; i++) {
        let sum = 0;
        // 计算输出层到隐藏层的梯度加权和
        for (let j = 0; j < this.outputSize; j++) {
          sum += this.Why[j][i] * dy[j];
        }
        // 结合下一个时间步的隐藏状态梯度和当前隐藏状态的导数，得到当前隐藏状态梯度
        dh[i] = (sum + dhNext[i]) * this.tanhDerivative(h[t][i]);
      }
      // 将当前输入转换为独热编码向量
      const x = this.oneHotEncode(inputs[t], this.inputSize);
      // 计算输入到隐藏层的权重梯度、隐藏层到隐藏层的权重梯度和隐藏层的偏置梯度
      for (let i = 0; i < this.hiddenSize; i++) {
        for (let j = 0; j < this.inputSize; j++) {
          // 累加当前时间步的输入到隐藏层的权重梯度
          dWxh[i][j] += dh[i] * x[j];
        }
        for (let j = 0; j < this.hiddenSize; j++) {
          // 累加当前时间步的隐藏层到隐藏层的权重梯度
          dWhh[i][j] += dh[i] * (t > 0 ? h[t - 1][j] : 0);
        }
        // 累加当前时间步的隐藏层的偏置梯度
        dbh[i] += dh[i];
      }
      // 更新下一个时间步的隐藏状态梯度为当前时间步的隐藏状态梯度
      dhNext = dh;
    }

    // 返回包含所有梯度的对象
    return { dWxh, dWhh, dWhy, dbh, dby };
  }

  // 生成指定行数和列数的零矩阵
  zerosMatrix(rows, cols) {
    const matrix = [];
    // 遍历每一行
    for (let i = 0; i < rows; i++) {
      // 生成指定长度的零向量并添加到矩阵中
      matrix.push(this.zeros(cols));
    }
    return matrix;
  }

  // 更新参数函数，根据计算得到的梯度更新权重和偏置
  updateParameters(dWxh, dWhh, dWhy, dbh, dby, learningRate) {
    // 更新输入到隐藏层的权重矩阵
    for (let i = 0; i < this.hiddenSize; i++) {
      for (let j = 0; j < this.inputSize; j++) {
        // 根据梯度和学习率更新权重
        this.Wxh[i][j] -= learningRate * dWxh[i][j];
      }
      // 更新隐藏层到隐藏层的权重矩阵
      for (let j = 0; j < this.hiddenSize; j++) {
        // 根据梯度和学习率更新权重
        this.Whh[i][j] -= learningRate * dWhh[i][j];
      }
    }
    // 更新隐藏层到输出层的权重矩阵和输出层的偏置向量
    for (let i = 0; i < this.outputSize; i++) {
      for (let j = 0; j < this.hiddenSize; j++) {
        // 根据梯度和学习率更新权重
        this.Why[i][j] -= learningRate * dWhy[i][j];
      }
      // 根据梯度和学习率更新偏置
      this.by[i] -= learningRate * dby[i];
    }
    // 更新隐藏层的偏置向量
    for (let i = 0; i < this.hiddenSize; i++) {
      // 根据梯度和学习率更新偏置
      this.bh[i] -= learningRate * dbh[i];
    }
  }

  // 保存模型参数到文件
  saveModel(filePath) {
    const modelData = {
      inputSize: this.inputSize,
      hiddenSize: this.hiddenSize,
      outputSize: this.outputSize,
      Wxh: this.Wxh,
      Whh: this.Whh,
      Why: this.Why,
      bh: this.bh,
      by: this.by
    };
    const jsonData = JSON.stringify(modelData);
    fs.writeFileSync(filePath, jsonData);
    console.log(`Model saved to ${filePath}`);
  }

  // 从文件加载模型参数
  static loadModel(filePath) {
    try {
      const jsonData = fs.readFileSync(filePath, 'utf8');
      const modelData = JSON.parse(jsonData);
      const rnn = new RNN(modelData.inputSize, modelData.hiddenSize, modelData.outputSize);
      rnn.Wxh = modelData.Wxh;
      rnn.Whh = modelData.Whh;
      rnn.Why = modelData.Why;
      rnn.bh = modelData.bh;
      rnn.by = modelData.by;
      console.log(`Model loaded from ${filePath}`);
      return rnn;
    } catch (error) {
      console.error('Error loading model:', error);
      return null;
    }
  }
}

// 问答函数，根据输入的问题生成回答
function askQuestion(question, rnn, wordToIndex, indexToWord) {
  // 将问题转换为小写，然后使用正则表达式匹配所有单词
  const questionWords = question.toLowerCase().match(/\w+/g);
  // 将问题中的每个单词转换为对应的索引，如果单词不在词汇表中则设为 0
  const inputSequence = questionWords.map(word => wordToIndex[word] || 0);
  // 对问题进行前向传播，得到输出
  const { y } = rnn.forward(inputSequence);
  // 获取最后一个时间步的输出
  const lastOutput = y[y.length - 1];
  // 找到输出中值最大的索引，即为预测的单词索引
  const predictedIndex = lastOutput.indexOf(Math.max(...lastOutput));
  // 根据预测的索引找到对应的单词
  const answer = indexToWord[predictedIndex];
  return answer;
}



async function main() {
  // 指定要读取的文本文件的路径，你需要将其替换为实际的文件路径
  const filePath = 'source.txt';
  // 使用 fs.readFileSync 同步读取文件内容，并以 UTF-8 编码解析，将结果存储在 text 变量中
  const text = fs.readFileSync(filePath, 'utf8');
  // 调用 preprocessText 函数对读取的文本进行预处理，并将结果解构赋值给相应的变量
  const { vocabulary, wordToIndex, indexToWord, inputSequence } = preprocessText(text);
  // return console.log('=================',{ vocabulary, wordToIndex, indexToWord, inputSequence } );

  let rnn;
  // 检查是否有保存的模型
  const modelFilePath = 'model.json';
  if (fs.existsSync(modelFilePath)) {
    console.log('Loading model from file...');
    rnn = RNN.loadModel(modelFilePath);
  } else {
    // 输入大小，即词汇表的长度
    const inputSize = vocabulary.length;
    // 隐藏层的大小
    const hiddenSize = 100;
    // 输出大小，即词汇表的长度
    const outputSize = vocabulary.length;
    // 创建 RNN 实例
    rnn = new RNN(inputSize, hiddenSize, outputSize);
    // 学习率，控制参数更新的步长
    const learningRate = 0.01;
    // 训练的轮数
    const numEpochs = 10;

    // 训练循环
    for (let epoch = 0; epoch < numEpochs; epoch++) {
      // 遍历输入序列，每次取一个输入和对应的目标
      for (let i = 0; i < inputSequence.length - 1; i++) {
        // 当前输入
        const inputs = [inputSequence[i]];
        // 对应的目标
        const targets = [inputSequence[i + 1]];
        // 进行前向传播，得到隐藏状态和输出
        const { h, y } = rnn.forward(inputs);
        // 进行反向传播，计算梯度
        const { dWxh, dWhh, dWhy, dbh, dby } = rnn.backward(inputs, targets, h, y);
        // 根据梯度更新参数
        rnn.updateParameters(dWxh, dWhh, dWhy, dbh, dby, learningRate);
      }
      // 输出当前训练轮数完成的信息
      console.log(`Epoch ${epoch + 1} completed`);
    }
    // 保存训练好的模型
    rnn.saveModel(modelFilePath);
  }

  // 示例问题，你可以修改为实际的问题
  const question = "my name is fang";
  // 调用 askQuestion 函数生成回答
  const answer = askQuestion(question, rnn, wordToIndex, indexToWord);
  // 输出问题
  console.log(`Question: ${question}`);
  // 输出回答
  console.log(`Answer: ${answer}`);
}

main();
